msg_tool\scripts\emote/
dref.rs

1//! Emote DPAK-referenced Image File (.dref)
2use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::img::*;
8use anyhow::Result;
9use emote_psb::PsbReader;
10use libtlg_rs::TlgColorType;
11use std::collections::HashMap;
12use std::io::Read;
13use std::path::{Path, PathBuf};
14use url::Url;
15
16#[derive(Debug)]
17/// Emote DREF Script Builder
18pub struct DrefBuilder {}
19
20impl DrefBuilder {
21    /// Creates a new instance of `DrefBuilder`
22    pub fn new() -> Self {
23        Self {}
24    }
25}
26
27impl ScriptBuilder for DrefBuilder {
28    fn default_encoding(&self) -> Encoding {
29        Encoding::Cp932
30    }
31
32    fn build_script(
33        &self,
34        buf: Vec<u8>,
35        filename: &str,
36        encoding: Encoding,
37        _archive_encoding: Encoding,
38        config: &ExtraConfig,
39        archive: Option<&Box<dyn Script>>,
40    ) -> Result<Box<dyn Script>> {
41        Ok(Box::new(Dref::new(
42            buf, encoding, filename, config, archive,
43        )?))
44    }
45
46    fn extensions(&self) -> &'static [&'static str] {
47        &["dref"]
48    }
49
50    fn script_type(&self) -> &'static ScriptType {
51        &ScriptType::EmoteDref
52    }
53
54    fn is_image(&self) -> bool {
55        true
56    }
57}
58
59struct Dpak {
60    psb: VirtualPsbFixed,
61}
62
63struct OffsetData {
64    left: u32,
65    top: u32,
66}
67
68impl Dpak {
69    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
70        let f = std::fs::File::open(path)?;
71        let mut f = std::io::BufReader::new(f);
72        let mut psb = PsbReader::open_psb(&mut f)
73            .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK: {:?}", e))?;
74        let psb = psb
75            .load()
76            .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK: {:?}", e))?;
77        let psb = psb.to_psb_fixed();
78        Ok(Self { psb })
79    }
80
81    pub fn load_from_data(data: &[u8]) -> Result<Self> {
82        let mut psb = PsbReader::open_psb(MemReaderRef::new(data))
83            .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK data: {:?}", e))?;
84        let psb = psb
85            .load()
86            .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK data: {:?}", e))?;
87        let psb = psb.to_psb_fixed();
88        Ok(Self { psb })
89    }
90
91    pub fn load_image(&self, name: &str) -> Result<(ImageData, Option<OffsetData>)> {
92        let root = self.psb.root();
93        let rid = root[name]
94            .resource_id()
95            .ok_or_else(|| anyhow::anyhow!("Resource ID for image '{}' not found in DPAK", name))?
96            as usize;
97        if rid >= self.psb.resources().len() {
98            return Err(anyhow::anyhow!(
99                "Resource ID {} out of bounds for DPAK with {} resources",
100                rid,
101                self.psb.resources().len()
102            ));
103        }
104        let resource = &self.psb.resources()[rid];
105        Self::load_img(&resource)
106    }
107
108    fn load_img(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
109        if libtlg_rs::is_valid_tlg(data) {
110            Ok((Self::load_tlg(data)?, None))
111        } else {
112            Self::load_png(data)
113        }
114    }
115
116    fn load_tlg(data: &[u8]) -> Result<ImageData> {
117        let img = libtlg_rs::load_tlg(MemReaderRef::new(data))
118            .map_err(|e| anyhow::anyhow!("Failed to decode TLG image: {:?}", e))?;
119        let color = img.color;
120        let mut re = ImageData {
121            width: img.width as u32,
122            height: img.height as u32,
123            color_type: match img.color {
124                TlgColorType::Grayscale8 => ImageColorType::Grayscale,
125                TlgColorType::Bgr24 => ImageColorType::Bgr,
126                TlgColorType::Bgra32 => ImageColorType::Bgra,
127            },
128            data: img.data,
129            depth: 8,
130        };
131        if let Some(v) = img.tags.get(&Vec::from(b"mode")) {
132            if v == b"alpha" && color == TlgColorType::Bgr24 {
133                convert_bgr_to_bgra(&mut re)?;
134            }
135        }
136        Ok(re)
137    }
138
139    fn load_png(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
140        let mut img = load_png(MemReaderRef::new(&data))?;
141        match img.color_type {
142            ImageColorType::Rgb => {
143                convert_rgb_to_rgba(&mut img)?;
144            }
145            _ => {}
146        }
147        Ok((
148            img,
149            Self::try_read_offset_from_png(MemReaderRef::new(&data))?,
150        ))
151    }
152
153    fn try_read_offset_from_png(mut data: MemReaderRef) -> Result<Option<OffsetData>> {
154        data.pos = 8; // Skip PNG signature
155        data.pos += 8; // Skip chunk size, type
156        data.pos += 17; // Skip IHDR chunk (length + type + width + height + bit depth + color type + compression method + filter method + interlace method)
157        loop {
158            let chunk_size = data.read_u32_be()?;
159            let mut chunk_type = [0u8; 4];
160            data.read_exact(&mut chunk_type)?;
161            if &chunk_type == b"IDAT" || &chunk_type == b"IEND" {
162                break;
163            }
164            if &chunk_type == b"oFFs" {
165                let x = data.read_u32_be()?;
166                let y = data.read_u32_be()?;
167                if data.read_u8()? == 0 {
168                    return Ok(Some(OffsetData { left: x, top: y }));
169                }
170            }
171            data.pos += chunk_size as usize + 4; // Skip chunk data and CRC
172        }
173        Ok(None)
174    }
175}
176
177#[derive(Default)]
178struct DpakLoader {
179    map: HashMap<String, Dpak>,
180}
181
182impl DpakLoader {
183    pub fn load_image(
184        &mut self,
185        dir: &Path,
186        dpak: &str,
187        filename: &str,
188    ) -> Result<(ImageData, Option<OffsetData>)> {
189        let dpak = match self.map.get(dpak) {
190            Some(d) => d,
191            None => {
192                let path = dir.join(dpak);
193                let ndpak = Dpak::new(&path)?;
194                self.map.insert(dpak.to_string(), ndpak);
195                self.map.get(dpak).unwrap()
196            }
197        };
198        dpak.load_image(filename)
199    }
200
201    pub fn load_archives(&mut self, in_archives: &HashMap<String, Vec<u8>>) -> Result<()> {
202        for (name, data) in in_archives.iter() {
203            if !self.map.contains_key(name) {
204                let dpak = Dpak::load_from_data(data)?;
205                self.map.insert(name.clone(), dpak);
206            }
207        }
208        Ok(())
209    }
210}
211
212/// Emote DREF Script
213pub struct Dref {
214    urls: Vec<Url>,
215    dir: PathBuf,
216    in_archives: HashMap<String, Vec<u8>>,
217}
218
219impl std::fmt::Debug for Dref {
220    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        f.debug_struct("Dref")
222            .field("urls", &self.urls)
223            .field("dir", &self.dir)
224            .finish()
225    }
226}
227
228impl Dref {
229    /// Create a new dref script
230    ///
231    /// * `buf` - The buffer containing the dref script data
232    /// * `encoding` - The encoding of the script
233    /// * `filename` - The name of the file
234    /// * `config` - Extra configuration options
235    /// * `archive` - Optional archive containing additional resources
236    pub fn new(
237        buf: Vec<u8>,
238        encoding: Encoding,
239        filename: &str,
240        _config: &ExtraConfig,
241        archive: Option<&Box<dyn Script>>,
242    ) -> Result<Self> {
243        let text = decode_with_bom_detect(encoding, &buf, true)?.0;
244        let mut urls = Vec::new();
245        for text in text.lines() {
246            let text = text.trim();
247            if text.is_empty() {
248                continue;
249            }
250            urls.push(Url::parse(text)?);
251        }
252        let path = Path::new(filename);
253        let dir = if let Some(parent) = path.parent() {
254            parent.to_path_buf()
255        } else {
256            PathBuf::from(".")
257        };
258        if urls.is_empty() {
259            return Err(anyhow::anyhow!("No URLs found in DREF file: {}", filename));
260        }
261        for u in urls.iter() {
262            if u.scheme() != "psb" {
263                return Err(anyhow::anyhow!(
264                    "Invalid URL scheme in DREF file: {} (expected 'psb')",
265                    u
266                ));
267            }
268        }
269        let mut in_archives = HashMap::new();
270        if let Some(archive) = archive {
271            if archive.is_archive() {
272                for url in urls.iter() {
273                    let filename = url.domain().ok_or(anyhow::anyhow!(
274                        "Invalid URL in DREF file: {} (missing domain)",
275                        url
276                    ))?;
277                    if let Ok(mut content) = archive.open_file_by_name(filename, true) {
278                        in_archives.insert(filename.to_string(), content.data()?);
279                    }
280                }
281            }
282        }
283        Ok(Self {
284            urls,
285            dir,
286            in_archives,
287        })
288    }
289}
290
291impl Script for Dref {
292    fn default_output_script_type(&self) -> OutputScriptType {
293        OutputScriptType::Json
294    }
295
296    fn default_format_type(&self) -> FormatOptions {
297        FormatOptions::None
298    }
299
300    fn is_image(&self) -> bool {
301        true
302    }
303
304    fn export_image(&self) -> Result<ImageData> {
305        let mut loader = DpakLoader::default();
306        loader.load_archives(&self.in_archives)?;
307        let base_url = &self.urls[0];
308        let dpak = base_url.domain().ok_or(anyhow::anyhow!(
309            "Invalid URL in DREF file: {} (missing domain)",
310            base_url
311        ))?;
312        let (mut base_img, base_offset) =
313            loader.load_image(&self.dir, dpak, base_url.path().trim_start_matches("/"))?;
314        if let Some(o) = base_offset {
315            eprintln!("WARN: Base image offset: left={}, top={}", o.left, o.top);
316            crate::COUNTER.inc_warning();
317        }
318        for url in &self.urls[1..] {
319            let dpak = url.domain().ok_or(anyhow::anyhow!(
320                "Invalid URL in DREF file: {} (missing domain)",
321                url
322            ))?;
323            let (mut img, img_offset) =
324                loader.load_image(&self.dir, dpak, url.path().trim_start_matches("/"))?;
325            let (top, left) = match img_offset {
326                Some(o) => (o.top, o.left),
327                None => (0, 0),
328            };
329            if base_img.color_type != img.color_type {
330                if base_img.color_type == ImageColorType::Rgba
331                    && img.color_type == ImageColorType::Rgb
332                {
333                    convert_rgb_to_rgba(&mut img)?;
334                } else if base_img.color_type == ImageColorType::Bgra
335                    && img.color_type == ImageColorType::Bgr
336                {
337                    convert_bgr_to_bgra(&mut img)?;
338                } else if base_img.color_type == ImageColorType::Rgba
339                    && img.color_type == ImageColorType::Bgra
340                {
341                    convert_bgra_to_rgba(&mut img)?;
342                } else if base_img.color_type == ImageColorType::Bgra
343                    && img.color_type == ImageColorType::Rgba
344                {
345                    convert_rgba_to_bgra(&mut img)?;
346                }
347            }
348            draw_on_img_with_opacity(&mut base_img, &img, left, top, 0xff)?;
349        }
350        Ok(base_img)
351    }
352}